import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.linear_model import LinearRegression
from sklearn.svm import SVR
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.neural_network import MLPRegressor
import joblib
import os

# 创建名为"model"的文件夹
os.makedirs("model", exist_ok=True)
# 通用的模型训练和评估函数

def train_and_evaluate_model(model, model_name, X_train, y_train, X_test, y_test, column):

    model.fit(X_train, y_train)  # 训练模型
    y_pred = model.predict(X_test)  # 预测
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    rems = np.sqrt(mse)
    # 保存模型到 "model" 文件夹中
    joblib.dump(model, f"model/{model_name}_{column}_model.pkl")

    return mse, r2, rems

# 单独的训练评估函数
def train_and_evaluate_lr(X_train, y_train, X_test, y_test, column):
    model = LinearRegression()
    return train_and_evaluate_model(model, 'Linear_Regression', X_train, y_train, X_test, y_test, column)

def train_and_evaluate_svr(X_train, y_train, X_test, y_test, column):
    model = SVR(kernel='rbf', C=1.0, epsilon=0.1)
    return train_and_evaluate_model(model, 'Support_Vector_Regression', X_train, y_train, X_test, y_test, column)

def train_and_evaluate_dtr(X_train, y_train, X_test, y_test, column):
    model = DecisionTreeRegressor(max_depth=None, min_samples_split=2, min_samples_leaf=1)
    return train_and_evaluate_model(model, 'Decision_Tree_Regression', X_train, y_train, X_test, y_test, column)

def train_and_evaluate_rfr(X_train, y_train, X_test, y_test, column):
    model = RandomForestRegressor(n_estimators=100, max_depth=None, min_samples_split=2, bootstrap=True, random_state=42)
    return train_and_evaluate_model(model, 'Random_Forest_Regression', X_train, y_train, X_test, y_test, column)

def train_and_evaluate_mlp(X_train, y_train, X_test, y_test, column):
    model = MLPRegressor(hidden_layer_sizes=(100,), max_iter=1000, activation='relu', solver='adam', random_state=42)
    return train_and_evaluate_model(model, 'MLP_Regression', X_train, y_train, X_test, y_test, column)

if __name__ == '__main__':
    # 存储结果
    results_data = []

    # 加载数据
    data = pd.read_csv('demo2.csv')

    # 设定目标变量y和特征变量X
    y = data.iloc[:, 0]
    X = data.iloc[:, 1:]

    # 获取特征列名列表
    columns = X.columns.tolist()

    for column in columns:
        # 根据每个特征变量进行训练和评估
        X_train, X_test, y_train, y_test = train_test_split(X[column].values.reshape(-1, 1), y, test_size=0.2, random_state=42)

        mse, r2, rems = train_and_evaluate_lr(X_train, y_train, X_test, y_test, column)
        results_data.append({'Model': 'Linear Regression', 'Feature': column, 'MSE': mse, 'R2': r2, 'REMS': rems})

        mse, r2, rems = train_and_evaluate_svr(X_train, y_train, X_test, y_test, column)
        results_data.append({'Model': 'Support Vector Regression', 'Feature': column, 'MSE': mse, 'R2': r2, 'REMS': rems})

        mse, r2, rems = train_and_evaluate_dtr(X_train, y_train, X_test, y_test, column)
        results_data.append({'Model': 'Decision Tree Regression', 'Feature': column, 'MSE': mse, 'R2': r2, 'REMS': rems})

        mse, r2, rems = train_and_evaluate_rfr(X_train, y_train, X_test, y_test, column)
        results_data.append({'Model': 'Random Forest Regression', 'Feature': column, 'MSE': mse, 'R2': r2, 'REMS': rems})

        mse, r2, rems = train_and_evaluate_mlp(X_train, y_train, X_test, y_test, column)
        results_data.append({'Model': 'MLP Regression', 'Feature': column, 'MSE': mse, 'R2': r2, 'REMS': rems})

    # 将结果保存到CSV文件
    results_df = pd.DataFrame(results_data)
    results_df.to_csv('model_evaluation_results.csv', index=False)